from flask import Flask, request, render_template, send_from_directory, jsonify
from PIL import Image
import os
import io
import uuid # 用于生成唯一文件名
import logging

app = Flask(__name__)
UPLOAD_FOLDER = 'uploads'
COMPRESSED_FOLDER = 'compressed'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['COMPRESSED_FOLDER'] = COMPRESSED_FOLDER

# 配置日志
logging.basicConfig(level=logging.INFO)
app.logger.setLevel(logging.INFO)


# 创建必要的文件夹
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(COMPRESSED_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_image():
    if 'file' not in request.files:
        app.logger.warning("Upload attempt with no file part.")
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        app.logger.warning("Upload attempt with no selected file.")
        return jsonify({'error': 'No selected file'}), 400

    if file and allowed_file(file.filename):
        original_filename = file.filename
        extension = original_filename.rsplit('.', 1)[1].lower()
        unique_filename_stem = str(uuid.uuid4())
        filename_on_server = f"{unique_filename_stem}.{extension}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename_on_server)

        try:
            file.save(filepath)
            app.logger.info(f"File '{original_filename}' uploaded as '{filename_on_server}'")
            return jsonify({
                'message': 'File uploaded successfully',
                'filename': filename_on_server,
                'original_client_filename': original_filename
            }), 200
        except Exception as e:
            app.logger.error(f"Error saving uploaded file '{original_filename}': {e}", exc_info=True)
            return jsonify({'error': f'Could not save file: {str(e)}'}), 500
    else:
        app.logger.warning(f"Upload attempt with disallowed file type: {file.filename}")
        return jsonify({'error': 'File type not allowed'}), 400

@app.route('/compress', methods=['POST'])
def compress_image_route():
    data = request.get_json()
    if not data:
        app.logger.warning("Compression attempt with no data.")
        return jsonify({'error': 'No data provided'}), 400

    original_server_filename = data.get('filename')
    original_client_filename = data.get('original_client_filename')  # 获取原始客户端文件名
    quality = int(data.get('quality', 75))

    if not original_server_filename:
        app.logger.warning("Compression attempt with no filename.")
        return jsonify({'error': 'No filename provided for compression'}), 400

    original_filepath = os.path.join(app.config['UPLOAD_FOLDER'], original_server_filename)
    if not os.path.exists(original_filepath):
        app.logger.error(f"Compression attempt: Original file '{original_server_filename}' not found.")
        return jsonify({'error': 'Original file not found on server'}), 404

    try:
        app.logger.info(f"Attempting to compress '{original_server_filename}' with quality {quality}")
        img = Image.open(original_filepath)
        original_format = img.format # 例如: 'JPEG', 'PNG'

        # 从原始客户端文件名中提取文件名部分（不含扩展名）
        if original_client_filename:
            base_filename, _ = os.path.splitext(original_client_filename)
        else:
            # 如果没有原始客户端文件名，使用服务器文件名
            base_filename, _ = os.path.splitext(original_server_filename)

        output_extension = 'jpg'
        output_format_pil = 'JPEG'
        save_options = {'quality': quality, 'optimize': True}

        # 根据原始格式确定输出格式和选项
        if original_format == 'WEBP':
            output_extension = 'webp'
            output_format_pil = 'WEBP'
            # WebP可以直接使用quality参数，Pillow会处理RGBA
            save_options = {'quality': quality}
            if img.mode == 'RGBA' and not img.info.get('icc_profile'): # 检查alpha但不是因为色彩配置文件
                 pass # Pillow的WebP保存器应该能处理RGBA
        elif original_format == 'PNG':
            # PNG转JPEG时处理透明度
            if img.mode == 'RGBA' or img.mode == 'P': # P模式可能有透明度
                img = img.convert('RGBA') # 确保是RGBA以访问alpha通道
                background = Image.new("RGB", img.size, (255, 255, 255)) # 白色背景
                background.paste(img, mask=img.split()[3]) # 粘贴时使用alpha通道作为mask
                img = background
            else: # 例如L模式或RGB模式的PNG
                img = img.convert('RGB') # 确保是RGB以保存为JPEG
        elif original_format == 'GIF':
            # GIF转JPEG时处理（仅第一帧）
            try:
                img.seek(0) # 定位到第一帧
                if img.mode == 'P' or img.mode == 'RGBA': # GIF常用P模式，也可能包含RGBA帧
                    # 如果有透明度，转换为RGBA再合并到RGB背景
                    if 'transparency' in img.info or img.mode == 'RGBA':
                        img_rgba = img.convert('RGBA')
                        background = Image.new("RGB", img_rgba.size, (255, 255, 255))
                        background.paste(img_rgba, mask=img_rgba.split()[3])
                        img = background
                    else:
                        img = img.convert('RGB')
            except EOFError: # 非动画GIF或只有一帧
                if img.mode == 'P' or img.mode == 'RGBA':
                    if 'transparency' in img.info or img.mode == 'RGBA':
                        img_rgba = img.convert('RGBA')
                        background = Image.new("RGB", img_rgba.size, (255, 255, 255))
                        background.paste(img_rgba, mask=img_rgba.split()[3])
                        img = background
                    else:
                        img = img.convert('RGB')
            if img.mode != 'RGB': # 确保最终是RGB
                 img = img.convert('RGB')
        else: # 其他格式（包括JPEG本身）默认输出JPEG
            if img.mode == 'RGBA' or img.mode == 'P':
                img = img.convert('RGBA')
                background = Image.new("RGB", img.size, (255, 255, 255))
                background.paste(img, mask=img.split()[3])
                img = background
            elif img.mode == 'CMYK': # 例如某些TIFF或JPEG可能是CMYK
                img = img.convert('RGB')
            elif img.mode != 'RGB': # 如果不是RGB，尝试转换
                img = img.convert('RGB')


        # 确保最终用于JPEG保存的图像是RGB模式
        if output_format_pil == 'JPEG' and img.mode != 'RGB':
            app.logger.info(f"Converting image from {img.mode} to RGB for JPEG output.")
            img = img.convert('RGB')

        # 使用原始文件名，只更改扩展名
        compressed_filename = f"{base_filename}.{output_extension}"
        compressed_filepath = os.path.join(app.config['COMPRESSED_FOLDER'], compressed_filename)

        img.save(compressed_filepath, format=output_format_pil, **save_options)
        app.logger.info(f"Image '{original_server_filename}' compressed to '{compressed_filename}'")

        return jsonify({
            'message': 'Image compressed successfully',
            'compressed_filename': compressed_filename,
            'download_url': f"/download/{compressed_filename}"
        }), 200

    except Exception as e:
        app.logger.error(f"Error during compression for '{original_server_filename}': {e}", exc_info=True)
        return jsonify({'error': f'Failed to compress image: {str(e)}'}), 500

@app.route('/download/<filename>')
def download_file(filename):
    return send_from_directory(app.config['COMPRESSED_FOLDER'], filename, as_attachment=True)

if __name__ == '__main__':
    app.run(debug=True) # debug=True 仅用于开发